#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
小红书数据管理系统 - 数据库操作工具
提供各种数据库常用操作函数，简化数据库交互
"""

import sqlite3
import os
import json
import datetime
from typing import Dict, List, Union, Any, Optional, Tuple
import threading

# 全局线程锁，防止多线程同时写入数据库
db_lock = threading.Lock()

# 数据库文件路径
DB_FILE = 'xhs_database.db'

def dict_factory(cursor, row):
    """将查询结果转换为字典格式"""
    d = {}
    for idx, col in enumerate(cursor.description):
        d[col[0]] = row[idx]
    return d

def get_db_connection():
    """获取数据库连接"""
    conn = sqlite3.connect(DB_FILE)
    conn.row_factory = dict_factory
    return conn

def execute_query(query: str, params: Tuple = ()) -> List[Dict]:
    """
    执行查询语句并返回结果
    
    Args:
        query: SQL查询语句
        params: 查询参数元组
    
    Returns:
        查询结果列表，每行数据为一个字典
    """
    conn = get_db_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(query, params)
        result = cursor.fetchall()
        return result
    finally:
        conn.close()

def execute_update(query: str, params: Tuple = ()) -> int:
    """
    执行更新语句并返回影响的行数
    
    Args:
        query: SQL更新语句
        params: 更新参数元组
    
    Returns:
        影响的行数
    """
    with db_lock:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(query, params)
            conn.commit()
            return cursor.rowcount
        finally:
            conn.close()

def execute_insert(query: str, params: Tuple = ()) -> int:
    """
    执行插入语句并返回新插入记录的ID
    
    Args:
        query: SQL插入语句
        params: 插入参数元组
    
    Returns:
        新插入记录的ID
    """
    with db_lock:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(query, params)
            conn.commit()
            return cursor.lastrowid
        finally:
            conn.close()

def execute_batch(queries: List[Dict[str, Any]]) -> bool:
    """
    批量执行SQL语句
    
    Args:
        queries: SQL语句列表，每个元素为一个字典，包含'query'和'params'两个键
    
    Returns:
        执行是否成功
    """
    with db_lock:
        conn = get_db_connection()
        try:
            cursor = conn.cursor()
            for query_dict in queries:
                cursor.execute(query_dict['query'], query_dict.get('params', ()))
            conn.commit()
            return True
        except Exception as e:
            conn.rollback()
            print(f"批量执行SQL失败: {e}")
            return False
        finally:
            conn.close()

def get_count(table: str, condition: str = '', params: Tuple = ()) -> int:
    """
    获取表中符合条件的记录数量
    
    Args:
        table: 表名
        condition: WHERE条件（不含WHERE关键字）
        params: 条件参数
    
    Returns:
        记录数量
    """
    query = f"SELECT COUNT(*) as count FROM {table}"
    if condition:
        query += f" WHERE {condition}"
    
    result = execute_query(query, params)
    if result:
        return result[0]['count']
    return 0

def get_records(table: str, 
                fields: str = '*', 
                condition: str = '', 
                params: Tuple = (),
                order_by: str = '',
                limit: int = 0,
                offset: int = 0) -> List[Dict]:
    """
    获取表中符合条件的记录
    
    Args:
        table: 表名
        fields: 要查询的字段，默认为所有字段
        condition: WHERE条件（不含WHERE关键字）
        params: 条件参数
        order_by: ORDER BY子句（不含ORDER BY关键字）
        limit: 限制返回的记录数
        offset: 跳过的记录数
    
    Returns:
        记录列表
    """
    query = f"SELECT {fields} FROM {table}"
    if condition:
        query += f" WHERE {condition}"
    if order_by:
        query += f" ORDER BY {order_by}"
    if limit:
        query += f" LIMIT {limit}"
    if offset:
        query += f" OFFSET {offset}"
    
    return execute_query(query, params)

def get_record_by_id(table: str, id_field: str, id_value: Any, fields: str = '*') -> Optional[Dict]:
    """
    根据ID获取表中的一条记录
    
    Args:
        table: 表名
        id_field: ID字段名
        id_value: ID值
        fields: 要查询的字段，默认为所有字段
    
    Returns:
        记录字典或None（如果记录不存在）
    """
    query = f"SELECT {fields} FROM {table} WHERE {id_field} = ?"
    result = execute_query(query, (id_value,))
    if result:
        return result[0]
    return None

def insert_record(table: str, record: Dict) -> int:
    """
    向表中插入一条记录
    
    Args:
        table: 表名
        record: 记录字典
    
    Returns:
        新插入记录的ID
    """
    fields = ', '.join(record.keys())
    placeholders = ', '.join(['?' for _ in record])
    query = f"INSERT INTO {table} ({fields}) VALUES ({placeholders})"
    return execute_insert(query, tuple(record.values()))

def update_record(table: str, id_field: str, id_value: Any, update_data: Dict) -> bool:
    """
    更新表中的一条记录
    
    Args:
        table: 表名
        id_field: ID字段名
        id_value: ID值
        update_data: 要更新的字段和值的字典
    
    Returns:
        更新是否成功
    """
    set_clause = ', '.join([f"{key} = ?" for key in update_data.keys()])
    query = f"UPDATE {table} SET {set_clause} WHERE {id_field} = ?"
    values = list(update_data.values())
    values.append(id_value)
    affected = execute_update(query, tuple(values))
    return affected > 0

def delete_record(table: str, id_field: str, id_value: Any) -> bool:
    """
    删除表中的一条记录
    
    Args:
        table: 表名
        id_field: ID字段名
        id_value: ID值
    
    Returns:
        删除是否成功
    """
    query = f"DELETE FROM {table} WHERE {id_field} = ?"
    affected = execute_update(query, (id_value,))
    return affected > 0

def export_db_to_json(file_path: str) -> bool:
    """
    将数据库导出到JSON文件
    
    Args:
        file_path: 导出文件路径
    
    Returns:
        导出是否成功
    """
    conn = get_db_connection()
    try:
        cursor = conn.cursor()
        # 获取所有表名
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
        tables = [row['name'] for row in cursor.fetchall()]
        
        # 导出每个表的数据
        export_data = {}
        for table in tables:
            cursor.execute(f"SELECT * FROM {table}")
            export_data[table] = cursor.fetchall()
        
        # 写入JSON文件
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(export_data, f, ensure_ascii=False, indent=2, default=str)
        
        return True
    except Exception as e:
        print(f"导出数据库到JSON失败: {e}")
        return False
    finally:
        conn.close()

def import_json_to_db(file_path: str) -> bool:
    """
    从JSON文件导入数据到数据库
    
    Args:
        file_path: JSON文件路径
    
    Returns:
        导入是否成功
    """
    if not os.path.exists(file_path):
        print(f"文件不存在: {file_path}")
        return False
    
    try:
        # 读取JSON文件
        with open(file_path, 'r', encoding='utf-8') as f:
            import_data = json.load(f)
        
        # 导入每个表的数据
        with db_lock:
            conn = get_db_connection()
            try:
                cursor = conn.cursor()
                for table, records in import_data.items():
                    # 清空表
                    cursor.execute(f"DELETE FROM {table}")
                    
                    # 插入记录
                    for record in records:
                        fields = ', '.join(record.keys())
                        placeholders = ', '.join(['?' for _ in record])
                        query = f"INSERT INTO {table} ({fields}) VALUES ({placeholders})"
                        cursor.execute(query, tuple(record.values()))
                
                conn.commit()
                return True
            except Exception as e:
                conn.rollback()
                print(f"导入JSON到数据库失败: {e}")
                return False
            finally:
                conn.close()
    except Exception as e:
        print(f"读取JSON文件失败: {e}")
        return False

def create_tables_from_schema(schema_file: str) -> bool:
    """
    从模式文件创建表
    
    Args:
        schema_file: SQL模式文件路径
    
    Returns:
        创建是否成功
    """
    if not os.path.exists(schema_file):
        print(f"模式文件不存在: {schema_file}")
        return False
    
    try:
        # 读取SQL文件
        with open(schema_file, 'r', encoding='utf-8') as f:
            schema_sql = f.read()
        
        # 分割SQL语句
        sql_statements = schema_sql.split(';')
        
        # 执行SQL语句
        with db_lock:
            conn = get_db_connection()
            try:
                cursor = conn.cursor()
                for statement in sql_statements:
                    statement = statement.strip()
                    if statement:
                        cursor.execute(statement)
                conn.commit()
                return True
            except Exception as e:
                conn.rollback()
                print(f"创建表失败: {e}")
                return False
            finally:
                conn.close()
    except Exception as e:
        print(f"读取模式文件失败: {e}")
        return False

def initialize_db_from_sql(sql_file: str) -> bool:
    """
    从SQL文件初始化数据库
    
    Args:
        sql_file: SQL文件路径
    
    Returns:
        初始化是否成功
    """
    if not os.path.exists(sql_file):
        print(f"SQL文件不存在: {sql_file}")
        return False
    
    try:
        # 读取SQL文件
        with open(sql_file, 'r', encoding='utf-8') as f:
            sql_content = f.read()
        
        # 分割SQL语句
        sql_statements = sql_content.split(';')
        
        # 执行SQL语句
        with db_lock:
            conn = get_db_connection()
            try:
                cursor = conn.cursor()
                for statement in sql_statements:
                    statement = statement.strip()
                    if statement:
                        try:
                            cursor.execute(statement)
                        except sqlite3.Error as e:
                            # 忽略特定错误
                            if "sqlite_sequence" in statement and "may not be dropped" in str(e):
                                print("忽略sqlite_sequence表删除错误")
                            elif "UNIQUE constraint failed" in str(e):
                                print(f"忽略唯一约束冲突错误: {e}")
                            else:
                                print(f"SQL错误: {e}")
                                print(f"问题语句: {statement[:100]}...")
                
                conn.commit()
                return True
            except Exception as e:
                conn.rollback()
                print(f"初始化数据库失败: {e}")
                return False
            finally:
                conn.close()
    except Exception as e:
        print(f"读取SQL文件失败: {e}")
        return False 